#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from retrying import retry
import requests
import sys
import os


@retry
def EfetchGeneXML(ids, parentpath, prefix, index):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi'
    payload = {"db": "gene", "id": ids, "retmode": "xml"}
    r = requests.post(url, stream=True, data=payload)
    if r.status_code == 200:
        with open(parentpath + "/" + prefix + "_" + index + ".xml", 'wb') as fd:
            for chunk in r.iter_content(1024 * 1024):
                fd.write(chunk)
            fd.close()
        if not Verify_Meta_XML(parentpath + "/" + prefix + "_" + index + ".xml"):
            print("Invaid XML. Try to redownload ...")
            raise IOError("Invalid xml.")
        else:
            print(prefix + "_" + index + ".xml downloaded.")
    else:
        print("Web Error. Retrying ...")
        raise IOError("Web Error.")


def Verify_Meta_XML(filename):
    import subprocess

    cmd = "grep '</Entrezgene-Set>' " + filename
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
    p.wait()
    result = ""
    for line in p.stdout.readlines():
        line = line.decode("utf-8").rstrip()
        result += line
    if result.find("</Entrezgene-Set>") != -1:
        # print(filename + " is valid.")
        return True
    elif result.find("No such file or directory") != -1:
        # print("Cannot find " + filename + ". Try to redownload...")
        return False
    else:
        # print(filename + " is invalid. Try to redownload...")
        return False


if __name__ == '__main__':
    parentpath = os.path.abspath(sys.argv[1])[:os.path.abspath(sys.argv[1]).rindex("/")]
    basename = os.path.basename(sys.argv[1])
    os.makedirs(parentpath + "/XMLs", exist_ok=True)
    prefix = basename[basename.index("_") + 1:basename.index(".")]

    infile = open(sys.argv[1], "r")
    infile.readline()

    index = 0
    count = 0
    ids = ""

    for line in infile:
        line = line.rstrip("\n")
        gene_id = line.split("\t")[1]
        ids += gene_id + ","
        count += 1
        if count == 5000:
            if os.path.exists(parentpath + "/XMLs/" + prefix + "_" + str(index).zfill(5) + ".xml"):
                print(prefix + "_" + str(index).zfill(5) + ".xml exists, skip.")
            else:
                EfetchGeneXML(ids[:-1], parentpath + "/XMLs", prefix, str(index).zfill(5))
            count = 0
            ids = ""
            index += 1
    if count > 0:
        if os.path.exists(parentpath + "/XMLs/" + prefix + "_" + str(index).zfill(5) + ".xml"):
            print(prefix + "_" + str(index).zfill(5) + ".xml exists, skip.")
        else:
            EfetchGeneXML(ids[:-1], parentpath + "/XMLs", prefix, str(index).zfill(5))
    infile.close()
